import os
from functools import partial
from multiprocessing import Pool
from pathlib import Path

import pandas as pd
from sklearn.model_selection import train_test_split
from tqdm import tqdm

from preprocessing.options import parse_args


def get_paths(image_name, train_crops_path):
    data = []
    image_path = os.path.join(train_crops_path, image_name)
    for image_crop in os.listdir(image_path):
        data.append(os.path.join(image_path, image_crop))
    return data


def save_csv(data_path, images, mode='train'):
    data = []
    with Pool(processes=1) as p:
        with tqdm(total=len(images)) as pbar:
            func = partial(get_paths, train_crops_path=data_path)
            for v in p.imap_unordered(func, images):
                pbar.update()
                data.extend(v)
    fold_data = []
    for img_path in data:
        path = Path(img_path)
        image = path.parent.name
        file = path.name
        fold_data.append([image, file])
    columns = ["image", "file"]
    pd.DataFrame(fold_data, columns=columns).to_csv('../data/data_{}.csv'.format(mode), index=False)


def save_test_csv(image_list):
    fold_data = []
    for image_name in image_list:
        if not image_name.endswith('.jpg'):
            continue
        fold_data.append(['test', image_name])
    columns = ["image", "file"]
    pd.DataFrame(fold_data, columns=columns).to_csv('../data/data_{}.csv'.format('test'), index=False)


def main():
    args = parse_args()
    train_crops_path = os.path.join(args.root_dir, 'train_crops')
    # train_mask_crops_path = os.path.join(args.root_dir, 'train_mask_crops')
    image_list = os.listdir(train_crops_path)
    print(len(image_list))
    images_train, images_val = train_test_split(image_list, test_size=0.1, random_state=111)
    # save_csv(train_crops_path, images_train, 'train')
    # save_csv(train_crops_path, images_val, 'val')

    test_path = os.path.join(args.root_dir, 'test')
    image_list = os.listdir(test_path)
    save_test_csv(image_list)


if __name__ == '__main__':
    main()
